import numpy as np
from sklearn.preprocessing import StandardScaler
from tensorflow.python.keras.models import Model
from tensorflow.python.keras import layers
from tensorflow.python.keras.layers import Input, Dense, LSTM
from sklearn.metrics import mean_absolute_error
from matplotlib import pyplot as plt
from tensorflow.python.keras.utils.vis_utils import plot_model
# from tensorflow.python.keras.utils import plot_model
import matplotlib
import matplotlib.pyplot as plt
import os




matplotlib.use('Agg')
plt.rcParams['font.sans-serif'] = ['Arial Black']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号


# 构建LSTM模型
def build_lstm(in_num, out_num, learn_rate=0.001, activate_func="sigmoid", loss_func="MAE"):
    layer_in = []
    concatenate = []
    for i in range(in_num):
        exec("in%s = Input(shape=(None, 2), name='input%s')" % (i, i))
        exec("x%s = LSTM(32, activation=activate_func, return_sequences=True)(in%s)" % (i, i))
        exec("x%s = LSTM(32, activation=activate_func)(x%s)" % (i, i))
        exec("concatenate.append(x%s)" % i)
        exec("layer_in.append(in%s)" % i)
    layer_out = []
    concatenated = layers.Concatenate()(concatenate)
    for i in range(out_num):
        exec("y%s = Dense(200)(concatenated)" % i)
        exec("layer_out.append(y%s)" % i)

    model1 = Model(layer_in, layer_out)
    model1.compile(optimizer='adam', loss=loss_func)

    return model1


def train(learn_rate=0.001, epoch=10, activate_func="sigmoid", loss_func="MAE"):
    datax_path = os.getcwd() + "/normalize_data/train/train_x.npy"
    datay_path = os.getcwd() + "/normalize_data/train/train_y.npy"
    # thruster_datax = np.load("./normalize_data/standardize/train_x.npy")
    thruster_datax = np.load(datax_path)
    in_shape = thruster_datax.shape[2]
    # thruster_datay = np.load("./normalize_data/standardize/train_y.npy")
    thruster_datay = np.load(datay_path)
    out_shape = thruster_datay.shape[2]
    lstm_model = build_lstm(out_shape, out_shape, learn_rate, activate_func, loss_func)
    train_data_x = []
    train_data_y = []
    val_data_x = []
    val_data_y = []
    for i in range(1, in_shape):
        exec(
            "train_data_x%s=np.concatenate((thruster_datax[:int(thruster_datax.shape[0]*0.8),:,0:1],thruster_datax[:int(thruster_datax.shape[0]*0.8),:,i:i+1]),axis=2)" % i)
        exec("train_data_y%s=thruster_datay[:int(thruster_datay.shape[0]*0.8),:,i-1:i]" % i)
        exec(
            "val_data_x%s=np.concatenate((thruster_datax[int(thruster_datax.shape[0]*0.8):,:,0:1],thruster_datax[int(thruster_datax.shape[0]*0.8):,:,i:i+1]),axis=2)" % i)
        exec("val_data_y%s=thruster_datay[int(thruster_datay.shape[0]*0.8):,:,i-1:i]" % i)

        xname = "train_data_x" + str(i)
        yname = "train_data_y" + str(i)

        valx = "val_data_x" + str(i)
        valy = "val_data_y" + str(i)

        train_data_x.append(locals()[xname])
        train_data_y.append(locals()[yname])
        val_data_x.append(locals()[valx])
        val_data_y.append(locals()[valy])

    history = lstm_model.fit(train_data_x, train_data_y, batch_size=5, epochs=epoch, verbose=1,
                             validation_data=(val_data_x, val_data_y))
    plt.figure(figsize=(5, 4), dpi=100)
    loss = history.history['loss']
    val_loss = history.history['val_loss']

    epochs = range(len(loss))

    plt.figure(figsize=(20, 16), dpi=80)
    plt.title(u'训练损失', fontproperties='SimHei', fontsize=40)
    plt.xlabel(u'训练轮次', fontproperties='SimHei', fontsize=40)
    plt.ylabel(u'损失值', fontproperties='SimHei', fontsize=40)
    plt.tick_params(labelsize=40)
    p1, = plt.plot(epochs, loss, 'bo', linewidth=10)
    p2, = plt.plot(epochs, val_loss, 'b', linewidth=10)
    # plt.title('Training and validation loss')
    plt.legend(handles=[p1, p2], prop={'family': 'SimHei', 'size': 40}, loc='best', labels=['训练集损失', '验证集损失'])

    plt.savefig('./img/loss.png')
    plt.close()

    lstm_model.save("./model/thruster_model.h5")
    data_path = "./model/thruster_model.h5"
    return data_path


if __name__ == '__main__':
    train()